package com.atguigu.bigdata.sparkSql

import org.apache.spark.SparkConf
import org.apache.spark.sql.{Row, SparkSession}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, LongType, StructType}

//自定义聚合函数
object SparkSQL04_UDAF {
  def main(args: Array[String]): Unit = {
    //sparkConf
    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("SparkSQL01")
    //sparkSession
    //    val session:SparkSession = new SparkSession(sparkConf)
    val spark = SparkSession.builder().config(sparkConf).getOrCreate()
//   创建聚合函数对象
    val udaf = new MyAgeAvgFunction1
//    注册聚合函数
    spark.udf.register("avgAge",udaf)
//    使用聚合函数
    var frame = spark.read.json("in/name.json")

    frame.createOrReplaceTempView("user")
    spark.sql("select age from user").show()
    spark.sql("select avgAge(age) from user").show()

    spark.stop()
  }

}
//  声明用户自定义函数
class MyAgeAvgFunction1 extends UserDefinedAggregateFunction{
// 函数输入的数据结构
  override def inputSchema: StructType = {
    new StructType().add("age",LongType)
  }
//计算时的数据结构
  override def bufferSchema: StructType = {
    new StructType().add("sum",LongType).add("count",LongType)
  }
//函数返回的数据类型
  override def dataType: DataType = DoubleType
//函数是否稳定
  override def deterministic: Boolean = true
//计算之前缓冲区的初始化
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
   buffer(0)=0L
   buffer(1)=0L
  }
//根据查询结果更新缓冲区数据
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
//    取到age的值进行累加-sum
      buffer(0) = buffer.getLong(0) + input.getLong(0)
//    count
      buffer(1) = buffer.getLong(1) + 1
  }
//将多个节点的缓冲区合并
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
//sum
    buffer1(0)=buffer1.getLong(0)+buffer2.getLong(0)
//count
    buffer1(1)=buffer1.getLong(1)+buffer2.getLong(1)

  }

  override def evaluate(buffer: Row): Any = {
    buffer.getLong(0).toDouble/buffer.getLong(1)
  }
}